import time
import logging
import numpy as np
import cv2

from ..utils import inference_engine as ie


class ImageOrientationCorrector(object):
    def __init__(self, model_path, mode=ie.MODE_ORT, res=(224,224),
                 need_norm=True, mean=(127.5, 127.5, 127.5), std=(127.5, 127.5, 127.5),
                 interpolation=cv2.INTER_AREA,
                 **kwargs):
        self.mode = mode
        self.res = res
        self.need_norm = need_norm
        self.mean = mean
        self.std = std
        self.interpolation = interpolation

        if mode == ie.MODE_ORT:
            self.channel_last = False
            self.engine = ie.ORTEngine(model_path, **kwargs)
        else:
            assert False, 'error: unsupported mode of inference engine.'

    def run(self, imgs):
        t0 = time.time()
        inputs = self.preprocess(imgs)
        t1 = time.time()
        outputs = self.inference(inputs)
        t2 = time.time()
        probs = self.postprocess(outputs)
        t3 = time.time()

        logging.debug('[ioc] pre: %.4fs | inf: %.4fs | post: %.4fs | total: %.4fs' %
                      (t1-t0, t2-t1, t3-t2, t3-t0))

        return probs

    def preprocess(self, imgs):
        t_imgs = [cv2.resize(img, self.res, interpolation=self.interpolation) for img in imgs]
        t_imgs = [self._transform(img) for img in t_imgs]
        inp = np.stack(t_imgs, axis=0)
        return [inp]

    def inference(self, inputs):
        outputs = self.engine.invoke(inputs)
        return outputs

    def postprocess(self, outputs):
        return outputs[0]

    def _transform(self, img):
        img = img.astype(np.float32)
        if self.need_norm:
            img = (img-np.float32(self.mean)) * (1/np.float32(self.std))
        if not self.channel_last:
            img = img.transpose(2, 0, 1).copy()
        return img
